//
// Created by yanhai on 2021/11/3.
//

// 线程安全hash
// 用在网络数据包寻找对应的flow的场景
// 使用模板，FlowType和PacketType类型必须具有指定的成员

#pragma once

#include <mutex>
#include <functional>
#include <cassert>

#include "utility/util/defer.h"

template<typename FlowType>
struct FlowHashBucket {

    // list
    FlowType *head;

    std::mutex mutex;
};

template<typename FlowType, typename PacketType>
class FlowHashTable {
public:
    FlowHashTable(unsigned int bucket_size,
                  const std::function<FlowType *(void)> &flow_allocator,
                  const std::function<void(FlowType *, const PacketType *)> &flow_init)
            : bucket_size_(bucket_size),
              flow_allocator_(flow_allocator),
              flow_init_(flow_init)
    {
        assert(bucket_size > 0);
        bucket_ = new FlowHashBucket<FlowType>[bucket_size_];
    }

    ~FlowHashTable()
    {
        delete[] bucket_;
    }

    // 通过packet获取所在的flow
    // 返回值要么是已经上锁的flow，要么是nullptr
    FlowType *Find(PacketType *p);

    // 通过packet获取所在的flow
    // 返回值要么是已经上锁的flow，要么是nullptr
    FlowType *FindAndNew(PacketType *p);

    bool Exist(const PacketType *p)
    {
        return Find(p) != nullptr;
    }

    // 桶的个数
    unsigned int BucketCount() const { return bucket_size_; }

    void Foreach(const std::function<bool(FlowType *f)> &callback);

private:
    bool FlowCompare(const FlowType *f, const PacketType *p) const
    {
        return *f == *p;
    }

private:
    FlowHashBucket<FlowType> *bucket_;
    unsigned int bucket_size_;
    const std::function<FlowType *(void)> &flow_allocator_;
    const std::function<void(FlowType *, const PacketType *)> &flow_init_;
};

template<typename FlowType, typename PacketType>
FlowType *FlowHashTable<FlowType, PacketType>::Find(PacketType *p)
{
    FlowType *f = nullptr;

    /* get our hash bucket and lock it */
    const uint32_t hash = p->flow_hash;
    FlowHashBucket<FlowType> *fb = &bucket_[hash % bucket_size_];

    std::lock_guard<std::mutex> lock(fb->mutex);

    /* see if the bucket already has a flow */
    if (fb->head == nullptr) {
        return nullptr;
    }

    /* ok, we have a flow in the bucket. Let's find out if it is our flow */
    f = fb->head;
    do {
        if (FlowCompare2(f, p) != 0) {
            f->mutex.lock();
            return f; /* return w/o releasing flow lock */
        }
        /* unless we removed 'f', prev_f needs to point to
         * current 'f' when adding a new flow below. */
        f = f->next;
    } while (f != nullptr);

    /* should be unreachable */
    return nullptr;
}


template<typename FlowType, typename PacketType>
FlowType *FlowHashTable<FlowType, PacketType>::FindAndNew(PacketType *p)
{
    FlowType *f = nullptr;

    /* get our hash bucket and lock it */
    const uint32_t hash = p->flow_hash;
    FlowHashBucket<FlowType> *fb = &bucket_[hash % bucket_size_];

    std::lock_guard<std::mutex> lock(fb->mutex);

    /* see if the bucket already has a flow */
    if (fb->head == nullptr) {
        f = flow_allocator_();
        if (f == nullptr) {
            return nullptr;
        }

        /* flow is locked */
        fb->head = f;

        /* got one, now lock, initialize and return */
        flow_init_(f, p);

        return f;
    }

    /* ok, we have a flow in the bucket. Let's find out if it is our flow */
    FlowType *prev_f = nullptr; /* previous flow */
    f = fb->head;
    do {
        FlowType *next_f = nullptr;
        if (FlowCompare(f, p)) {
            f->mutex.lock();
            return f; /* return w/o releasing flow lock */
        }
        /* unless we removed 'f', prev_f needs to point to
         * current 'f' when adding a new flow below. */
        prev_f = f;
        next_f = f->next;

        if (next_f == nullptr) {
            f = flow_allocator_();
            if (f == nullptr) {
                return nullptr;
            }

            /* flow is locked */

            // 头插法
            f->next = fb->head;
            fb->head = f;

            /* initialize and return */
            flow_init_(f, p);
            return f;
        }
        f = next_f;
    } while (f != nullptr);

    /* should be unreachable */
    return nullptr;
}

template<typename FlowType, typename PacketType>
void FlowHashTable<FlowType, PacketType>::Foreach(const std::function<bool(FlowType *f)> &callback)
{
    for (unsigned int idx = 0; idx < bucket_size_; idx++) {
        FlowHashBucket<FlowType> *fb = &bucket_[idx];

        // 尝试加锁
        if (!fb->mutex.try_lock()) {
            std::cout << "FlowHashForeach: fb try lock failed." << std::endl;
            continue;
        }

        // 加锁成功，延迟解锁
        defer BucketUnlock([fb]() { fb->mutex.unlock(); });

        FlowType *f = fb->head;
        while (f != nullptr) {

            // try lock
            if (!f->mutex.try_lock()) {
                std::cout << "FlowHashForeach: flow try lock failed." << std::endl;
                f = f->next;
                continue;
            }

            // unlock
            defer FlowUnlock([f]() { f->mutex.unlock(); });

            if (!callback(f))
                break;

            f = f->next;
        }
    }
}
